#include "system_hf.h"
#include "help.h"
#include <algorithm>
#include <cmath>
#include <fstream>
using std::abs;
using std::max;
using std::min;

void System_Tz_HF::setupOrbitals() {
  hfsolver.init();
  hfsolver.setbeta(beta);
  hfsolver.iter();
  hfsolver.calEnergy();
  HFenergy = hfsolver.Energy;
  Orbitals = hfsolver.HFOrbitals;
  stable_sort(Orbitals.begin(), Orbitals.end(),
              [](const OrbitalType &orb1, const OrbitalType &orb2) {
                return orb1.state < orb2.state;
              });
  setupGroups();
  // set fermi surface
  FermiSurface = 0;
  for (int i = 0; i < Orbitals.size(); i++) {
    if (Orbitals[i].state)
      break;
    FermiSurface = i;
  }
}

/*
  dim of MatEle=4,the 4 double are:
      0  interaction
      1  (p1 \dot p2)/(m hbar omega) + m omega r1 \dot r2 / hbar
      2  (r1-r2)^2 m omega/(2 hbar)
      3  (p1 \dot p2)/(m hbar omega)
*/
const int System_Tz_HF::MatEledim = System_TzType::MatEledim;
void System_Tz_HF::setupTwoBodyMat() {
  setupIndiceIn2BStates();
  int totalChannels = TwoBodyChannels.size();
  TwoBodyMat.resize(totalChannels);
#pragma omp parallel for
  for (int channel = 0; channel < totalChannels; channel++) {
    int num2BInChannel = TwoBodyChannels[channel].size();
    TwoBodyMat[channel].clear();
    TwoBodyMat[channel].resize((num2BInChannel * (num2BInChannel + 1)) / 2);
    for (int bra = 0; bra < num2BInChannel; bra++)
      for (int ket = bra; ket < num2BInChannel; ket++) {
        int ab = TwoBodyChannels[channel][bra];
        int cd = TwoBodyChannels[channel][ket];
        int a = TwoBodyStates[ab].a;
        int b = TwoBodyStates[ab].b;
        int c = TwoBodyStates[cd].a;
        int d = TwoBodyStates[cd].b;
        int J = TwoBodyStates[ab].J;
        TwoBodyMatEle MatEle(MatEledim);
        get2BmatOnNeed(a, b, c, d, J, MatEle);
        set2Bmat(channel, bra, ket, MatEle);
      }
  }
}

void System_Tz_HF::get2BmatOnNeed(int a, int b, int c, int d, int J,
                                  TwoBodyMatEle &MatEle) const {
  MatEle.resize(MatEledim);
  MatEle.setZero();
  int jja = Orbitals[a].jj;
  int jjb = Orbitals[b].jj;
  int jjc = Orbitals[c].jj;
  int jjd = Orbitals[d].jj;
  int Tzab = (Orbitals[a].tz + Orbitals[b].tz) / 2;
  int Tzcd = (Orbitals[c].tz + Orbitals[d].tz) / 2;

  if (a == b && J % 2)
    return;
  if (c == d && J % 2)
    return;
  if (Tzab != Tzcd)
    return; // isospin projection check
  if ((Orbitals[a].l + Orbitals[b].l + Orbitals[c].l + Orbitals[d].l) % 2)
    return; // parity check
  int Jmin = max(abs(jja - jjb) / 2, abs(jjc - jjd) / 2);
  int Jmax = min((jja + jjb) / 2, (jjc + jjd) / 2);
  if (Jmin > Jmax)
    return; // angular momenta check

  for (int ia = 0; ia < Orbitals[a].Coeff.size(); ia++)
    for (int ib = 0; ib < Orbitals[b].Coeff.size(); ib++)
      for (int ic = 0; ic < Orbitals[c].Coeff.size(); ic++)
        for (int id = 0; id < Orbitals[d].Coeff.size(); id++) {
          System_TzType::TwoBodyMatEle Temp;
          int iia = pSystem_Tz->Groups[Orbitals[a].GroupIndex][ia];
          int iib = pSystem_Tz->Groups[Orbitals[b].GroupIndex][ib];
          int iic = pSystem_Tz->Groups[Orbitals[c].GroupIndex][ic];
          int iid = pSystem_Tz->Groups[Orbitals[d].GroupIndex][id];
          pSystem_Tz->get2BmatAt(iia, iib, iic, iid, J, Temp);
          if (Temp.empty())
            continue;
          if (iia == iib)
            Temp *= sqrt(2.);
          if (iic == iid)
            Temp *= sqrt(2.);
          MatEle += Temp * (Orbitals[a].Coeff[ia] * Orbitals[b].Coeff[ib] *
                            Orbitals[c].Coeff[ic] * Orbitals[d].Coeff[id]);
        }
  if (a == b)
    MatEle *= sqrt(2.) / 2;
  if (c == d)
    MatEle *= sqrt(2.) / 2;
}

// void System_Tz_HF::setupTwoBodyMat()
// {
//    setupIndiceIn2BStates();
//   int totalChannels=TwoBodyChannels.size();
//   TwoBodyMat.resize(totalChannels);
//   for(int channel=0;channel<totalChannels;channel++)
//     {
//       int num2BInChannel=TwoBodyChannels[channel].size();
// TwoBodyMat[channel].clear();
//       TwoBodyMat[channel].resize( (num2BInChannel*(num2BInChannel+1))/2 );

//       int Vdim=pSystem_Tz->TwoBodyChannels[channel].size();
//       Mat2B Vmat(Vdim,Vdim);
//       //setup Vmat
//       for(int i=0;i<Vdim;i++)
// 	{
// 	  for(int j=i;j<Vdim;j++)
// 	    {
// 	      pSystem_Tz->get2Bmat(channel,i,j,Vmat(i,j));
// 	      Vmat(j,i)=Vmat(i,j);
// 	    }
// 	}
//       //setup TransMat
//       Mat TransMat(Vdim,num2BInChannel);
//       for(int i=0;i<Vdim;i++)
// 	for(int j=0;j<num2BInChannel;j++)
// 	  {
// 	    int ab=pSystem_Tz->TwoBodyChannels[channel][i];
// 	    int alphabeta=TwoBodyChannels[channel][j];
// 	    int a=pSystem_Tz->TwoBodyStates[ab].a;
// 	    int b=pSystem_Tz->TwoBodyStates[ab].b;
// 	    int alpha=TwoBodyStates[alphabeta].a;
// 	    int beta=TwoBodyStates[alphabeta].b;
// 	    //debug: TwoBodyState ab and alphabeta should have same quantum
// number 	    DataType coef=overlap(a,alpha)*overlap(b,beta); 	    int
// jjalpha=Orbitals[alpha].jj; 	    int jjbeta=Orbitals[beta].jj; 	    int
// J=TwoBodyStates[alphabeta].J; 	    coef-=phase( (jjalpha+jjbeta)/2 -J
// )*overlap(a,beta)*overlap(b,alpha); 	    if(a==b) coef/=sqrt(2.);
// if(alpha==beta) coef/=sqrt(2.); 	    TransMat(i,j)=coef;
// 	  }
//       cout<<channel<<endl;
//       //transform  and write into TwoBodyMat
//       for(int i=0;i<num2BInChannel;i++)
//       	for(int j=i;j<num2BInChannel;j++)
//       	  {
//       	    TwoBodyMatEle VHF(MatEledim);
//       	    for(int k=0;k<Vdim;k++)
//       	      for(int l=0;l<Vdim;l++)
//       		{
// 		  if(TransMat(l,j)==0||TransMat(k,i)==0) continue;
// 		  if(Vmat(k,l).empty()) continue;
//       		  VHF+=Vmat(k,l)*(TransMat(l,j)*TransMat(k,i));
//       		}
//       	    set2Bmat(channel,i,j,VHF);
//       	  }
//     }
// }
System_Tz_HF::DataType System_Tz_HF::get1B(int bra, int ket) const {
  if (Flag_TBME == 2) {
    // Only for NNLO sat
    if (Orbitals[bra].isInAGroupWith(Orbitals[ket])) {
      int group = GroupIndice[bra];
      int a = OrbIndice[bra];
      int b = OrbIndice[ket];
      int dim = Groups[group].size();
      if (a > b)
        swap(a, b);
      return OneBodyMat[group][mapab_e(a, b, dim)];
    } else
      return 0;
  } else {
    DataType val = 0;
    return (1 - 1. / realA) * get1B_p2(bra, ket);
  }
}

System_Tz_HF::DataType System_Tz_HF::get1B_p2(int bra, int ket) const {
  DataType val = 0;
  if (Orbitals[bra].isInAGroupWith(Orbitals[ket])) {
    for (int ibra = 0; ibra < Orbitals[bra].Coeff.size(); ibra++)
      for (int iket = 0; iket < Orbitals[ket].Coeff.size(); iket++) {
        int iibra = pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
        int iiket = pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
        val += pSystem_Tz->get1B_p2(iibra, iiket) * Orbitals[bra].Coeff[ibra] *
               Orbitals[ket].Coeff[iket];
      }
  }
  return val;
}
System_Tz_HF::DataType System_Tz_HF::get1B_p2_plus_r2(int bra, int ket) const {
  DataType val = 0;
  if (Orbitals[bra].isInAGroupWith(Orbitals[ket])) {
    for (int ibra = 0; ibra < Orbitals[bra].Coeff.size(); ibra++)
      for (int iket = 0; iket < Orbitals[ket].Coeff.size(); iket++) {
        int iibra = pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
        int iiket = pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
        if (iibra != iiket)
          continue;
        val += pSystem_Tz->Orbitals[iibra].e * Orbitals[bra].Coeff[ibra] *
               Orbitals[ket].Coeff[iket];
        // val+=pSystem_Tz->get1B_p2_plus_r2(iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
      }
  }
  return val;
}

System_Tz_HF::DataType System_Tz_HF::rL(int bra, int ket, int L) const {
  DataType val = 0;
  for (int ibra = 0; ibra < Orbitals[bra].Coeff.size(); ibra++)
    for (int iket = 0; iket < Orbitals[ket].Coeff.size(); iket++) {
      int iibra = pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
      int iiket = pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
      val += pSystem_Tz->rL(iibra, iiket, L) * Orbitals[bra].Coeff[ibra] *
             Orbitals[ket].Coeff[iket];
    }
  return val;
}
System_Tz_HF::DataType System_Tz_HF::Q(int p, int h, int L, int isospin) const {
  int Tzh = Orbitals[h].tz;
  int Tzp = Orbitals[p].tz;
  if (Tzh != Tzp)
    return 0;
  int jjp = Orbitals[p].jj;
  int jjh = Orbitals[h].jj;
  int lp = Orbitals[p].l;
  int lh = Orbitals[h].l;
  DataType temp, fac;
  if (L == 0) {
    temp = rL(p, h, 2);
    temp *= reducedYMatEle(lp, jjp, lh, jjh, 0) * sqrt(4 * Pi);
  } else {
    temp = rL(p, h, L);
    temp *= reducedYMatEle(lp, jjp, lh, jjh, L);
  }
  fac = isospin ? Tzh : 1.0;
  temp *= fac;
  return temp;
}

System_Tz_HF::DataType System_Tz_HF::Q1_v(int p, int h) const {
  int Tzh = Orbitals[h].tz;
  int Tzp = Orbitals[p].tz;
  if (Tzh != Tzp)
    return 0;
  DataType temp = rL(p, h, 1);
  int jjp = Orbitals[p].jj;
  int jjh = Orbitals[h].jj;
  int lp = Orbitals[p].l;
  int lh = Orbitals[h].l;
  temp *= reducedYMatEle(lp, jjp, lh, jjh, 1);
  temp *= (Tzh * 0.5 - 0.5 + Z * 1.0 / realA);
  return temp;
}
System_Tz_HF::DataType System_Tz_HF::Q1_s(int p, int h, DataType Rms) const {
  int Tzh = Orbitals[h].tz;
  int Tzp = Orbitals[p].tz;
  if (Tzh != Tzp)
    return 0;
  DataType temp = rL(p, h, 3) - rL(p, h, 1) * Rms * Rms * 5. / 3;

  int jjp = Orbitals[p].jj;
  int jjh = Orbitals[h].jj;
  int lp = Orbitals[p].l;
  int lh = Orbitals[h].l;
  temp *= reducedYMatEle(lp, jjp, lh, jjh, 1);
  return temp;
}

System_Tz_HF::DataType System_Tz_HF::Beta_F(int type, int a, int b) const {
  DataType temp(0);
  int Tza = Orbitals[a].tz;
  int Tzb = Orbitals[b].tz;
  int jja = Orbitals[a].jj;
  int jjb = Orbitals[b].jj;
  int la = Orbitals[a].l;
  int lb = Orbitals[b].l;
  if ((jja != jjb) || (la != lb))
    return 0;
  if (Tza - Tzb == 2 * type) {
    temp = radialOverlap(a, b);
  }
  temp *= sqrt(jja + 1.);
  return temp;
}
System_Tz_HF::DataType System_Tz_HF::radialOverlap(int a, int b) const {
  //!!works only with same neutron and proton orbital number as well as good
  //! order for
  // the HO. basis, the dot product of the coefficients vectors is used as the
  // overlap. This works only if the underlying HO basis is in the same order.
  int dim_a = Orbitals[a].Coeff.size();
  int dim_b = Orbitals[b].Coeff.size();
  if (dim_a != dim_b) {
    std::cerr << "error in radialOverlap in system_hf\n";
    exit(-1);
  }
  DataType temp(0);
  for (int i = 0; i < dim_a; i++) {
    temp += Orbitals[a].Coeff[i] * Orbitals[b].Coeff[i];
  }
  return temp;
}
System_Tz_HF::DataType System_Tz_HF::Beta_GT(int type, int a, int b) const {
  DataType temp(0);
  int Tza = Orbitals[a].tz;
  int Tzb = Orbitals[b].tz;
  int jja = Orbitals[a].jj;
  int jjb = Orbitals[b].jj;
  int la = Orbitals[a].l;
  int lb = Orbitals[b].l;
  if ((la != lb) || (jja - jjb) > 2 || (jja + jjb) < 2)
    return 0;
  if (Tza - Tzb == 2 * type) {
    temp = radialOverlap(a, b);
    temp *= 2 * reducedSMatEle(la, jja, lb, jjb);
  }
  return temp;
}

void System_Tz_HF::setupFromFile(const char *orbfile, const char *intfile) {
  std::ifstream orbfin(orbfile);
  double trash;
  int num;
  orbfin >> num;
  int l, jj, tz;
  double ereal, eimag;
  Orbitals.clear();
  Orbitals.reserve(num);
  int line, ntrash;

  for (int i = 0; i < num; i++) {
    orbfin >> line >> ntrash >> ntrash >> ntrash >> l >> jj >> tz;
    orbfin >> ereal >> eimag;
    Orbitals.push_back(System_Tz_HF::OrbitalType(l, jj, tz, ereal));
  }
  setupGroups();

  std::ifstream vfin(intfile);
  setupTwoBodyStates();
  setupIndiceIn2BStates();
  int totalChannels = TwoBodyChannels.size();
  TwoBodyMat.resize(totalChannels);
  for (int channel = 0; channel < totalChannels; channel++) {
    int num2BInChannel = TwoBodyChannels[channel].size();
    TwoBodyMat[channel].clear();
    TwoBodyMat[channel].resize((num2BInChannel * (num2BInChannel + 1)) / 2);
  }
  int totalnum;
  int J;
  int a, b, c, d;
  double vreal, vimag;
  vfin >> totalnum;
  for (int i = 0; i < totalnum; i++) {
    vfin >> a >> b >> c >> d >> J >> vreal >> vimag;
    a--;
    b--;
    c--;
    d--;
    System_Tz_HF::TwoBodyMatEle MatEle(System_Tz_HF::MatEledim);
    MatEle[0] = vreal;
    set2BmatAt(a, b, c, d, J, MatEle);
  }
  FermiSurface = 5;
}

void setup_nnlo_sat(System_Tz_HO &system_ho, System_Tz_HF &system_hf,
                    const char *spfile, const char *umatfile,
                    const char *kinfile, const char *vfile) {
  using std::ifstream;
  ifstream spfin(spfile);
  ifstream umatfin(umatfile);
  ifstream vfin(vfile);
  ifstream kinfin(kinfile);
  // setup System_Tz_HO orbitals and System_Tz_HF orbitals, we use spfile both
  // for HO and HF, because HO orbitals has same structure with HF. orbitals
  double trash;
  spfin >> trash;
  int line;
  int n, l, jj, tz;
  double ereal, eimag;
  system_ho.Orbitals.clear();
  system_hf.Orbitals.clear();
  while (spfin >> line) {
    spfin >> n >> l >> jj >> tz;
    spfin >> ereal >> eimag;
    system_ho.Orbitals.push_back(System_Tz_HO::OrbitalType(n, l, jj, tz, 0));
    system_hf.Orbitals.push_back(System_Tz_HF::OrbitalType(l, jj, tz, ereal));
  }
  system_ho.setupGroups();
  system_hf.setupGroups();
  double hbar_omega = 22;
  system_ho.hbar_omega = hbar_omega;
  system_ho.lengthb = hbar_c / sqrt(hbar_omega * mc2);
  system_hf.hbar_omega = hbar_omega;
  // setup System_Tz_HF
  int a, b;

  for (int i = 0; i < system_hf.Orbitals.size(); i++) {
    int group = system_ho.GroupIndice[i];
    system_hf.Orbitals[i].GroupIndex = group;
    int dim = system_ho.Groups[group].size();
    system_hf.Orbitals[i].Coeff.resize(dim);
  }

  double coeff, temp;
  for (int i = 0; i < system_hf.Orbitals.size(); i++) {
    for (int j = 0; j < system_hf.Orbitals.size(); j++) {
      umatfin >> a >> b >> coeff >> temp;
      a--;
      b--;
      if (system_hf.Orbitals[a].isInAGroupWith(system_hf.Orbitals[b])) {
        int index = system_ho.OrbIndice[a];
        system_hf.Orbitals[j].Coeff[index] = coeff;
      }
    }
  }
  vector<int> indexs(system_hf.Orbitals.size());
  for (int i = 0; i < indexs.size(); i++)
    indexs[i] = i;
  sort(indexs.begin(), indexs.end(), [&](int i1, int i2) {
    return system_hf.Orbitals[i1] < system_hf.Orbitals[i2];
  });
  vector<int> orders(indexs.size());
  for (int i = 0; i < indexs.size(); i++) {
    orders[indexs[i]] = i;
  }
  sort(system_hf.Orbitals.begin(), system_hf.Orbitals.end());

  vector<System_Tz_HF::OrbitalType> &HFOrbitals = system_hf.Orbitals;

  // fill
  int resZ = system_hf.Z;
  int resN = system_hf.A - system_hf.Z;
  for (int i = 0; i < HFOrbitals.size(); i++) {
    // fill proton
    if (HFOrbitals[i].tz == -1) {
      int num = HFOrbitals[i].jj + 1;
      resZ -= num;
      if (resZ >= 0) {
        HFOrbitals[i].state = 0; // hole
        HFOrbitals[i].OccuP = 1.0;
      } else if ((resZ + num) > 0) {
        HFOrbitals[i].state = 0; // hole, but not fully occupied.
        HFOrbitals[i].OccuP = double(resZ + num) / num;
      } else {
        HFOrbitals[i].state = 1; // particle
        HFOrbitals[i].OccuP = 0.0;
      }
    }
    // fill neutron
    if (HFOrbitals[i].tz == 1) {
      int num = HFOrbitals[i].jj + 1;
      resN -= num;
      if (resN >= 0) {
        HFOrbitals[i].state = 0; // hole
        HFOrbitals[i].OccuP = 1.0;
      } else if ((resN + num) > 0) {
        HFOrbitals[i].state = 0; // hole, but not fully occupied.
        HFOrbitals[i].OccuP = double(resN + num) / num;
      } else {
        HFOrbitals[i].state = 1; // particle
        HFOrbitals[i].OccuP = 0.0;
      }
    }
  }

  stable_sort(indexs.begin(), indexs.end(), [&](int i1, int i2) {
    return system_hf.Orbitals[orders[i1]].state <
           system_hf.Orbitals[orders[i2]].state;
  });
  for (int i = 0; i < indexs.size(); i++) {
    orders[indexs[i]] = i;
  }
  stable_sort(system_hf.Orbitals.begin(), system_hf.Orbitals.end(),
              [](const System_Tz_HF::OrbitalType &orb1,
                 const System_Tz_HF::OrbitalType &orb2) {
                return orb1.state < orb2.state;
              });
  system_hf.setupGroups();

  // set fermi surface
  int Atemp = system_hf.A;
  for (int i = 0; i < system_hf.Orbitals.size(); i++) {
    Atemp -= system_hf.Orbitals[i].jj + 1;
    if (Atemp < 0)
      break;
    system_hf.FermiSurface = i;
  }
  // setup 1Bmat of System_Tz_HF
  system_hf.OneBodyMat.resize(system_hf.Groups.size());
  for (int i = 0; i < system_hf.Groups.size(); i++) {
    int dim = system_hf.Groups[i].size();
    system_hf.OneBodyMat[i].resize((dim * (dim + 1)) / 2);
  }
  double kin;
  for (int i = 0; i < system_hf.Orbitals.size(); i++) {
    for (int j = 0; j < system_hf.Orbitals.size(); j++) {
      kinfin >> a >> b >> kin >> temp >> temp >> temp;
      a--;
      b--;
      a = orders[a];
      b = orders[b];
      if (system_hf.Orbitals[a].isInAGroupWith(system_hf.Orbitals[b]) &&
          a <= b) {
        int group = system_hf.GroupIndice[a];
        int dim = system_hf.Groups[group].size();
        int aa = system_hf.OrbIndice[a];
        int bb = system_hf.OrbIndice[b];
        system_hf.OneBodyMat[group][mapab_e(aa, bb, dim)] = kin;
      }
    }
  }
  // setup 2Bmat of System_Tz_HF
  system_hf.setupTwoBodyStates();
  system_hf.setupIndiceIn2BStates();
  int totalChannels = system_hf.TwoBodyChannels.size();
  system_hf.TwoBodyMat.resize(totalChannels);
  for (int channel = 0; channel < totalChannels; channel++) {
    int num2BInChannel = system_hf.TwoBodyChannels[channel].size();
    system_hf.TwoBodyMat[channel].clear();
    system_hf.TwoBodyMat[channel].resize(
        (num2BInChannel * (num2BInChannel + 1)) / 2);
  }
  int Tz, Par, J;
  int c, d;
  double v;
  while (vfin >> Tz) {
    vfin >> Par >> J >> a >> b >> c >> d >> v;
    a--;
    b--;
    c--;
    d--;
    a = orders[a];
    b = orders[b];
    c = orders[c];
    d = orders[d];
    System_Tz_HF::TwoBodyMatEle MatEle(System_Tz_HF::MatEledim);
    MatEle[0] = v;
    system_hf.set2BmatAt(a, b, c, d, J, MatEle);
  }
}

double nnlo_sat_getHFenergy(
    const System_Tz_HF &system_hf) { // for nnlo ,Ehf=1/6(2*t + 4 sp - twobody)
  double E = 0;
  for (int i = 0; i <= system_hf.FermiSurface; i++) {
    E += 2 * system_hf.get1B(i, i) * (system_hf.Orbitals[i].jj + 1);
    E += 4 * system_hf.Orbitals[i].e * (system_hf.Orbitals[i].jj + 1);
    for (int j = 0; j <= system_hf.FermiSurface; j++) {
      int JJ1 = system_hf.Orbitals[i].jj;
      int JJ2 = system_hf.Orbitals[j].jj;
      int Jmin = abs(JJ1 - JJ2) / 2;
      int Jmax = (JJ1 + JJ2) / 2;
      for (int J = Jmin; J <= Jmax; J++) {
        E -= system_hf.get2B(i, j, i, j, J) * (2 * J + 1);
      }
    }
  }
  return E / 6;
}
